/**
 * Copyright 2009 Erick Armbrust 
 * Licensed under the Apache License, Version 2.0 (the "License"); 
 * you may not use this file except in compliance with the License. 
 * You may obtain a copy of the License at 
 * 
 *   http://www.apache.org/licenses/LICENSE-2.0 
 *   
 * Unless required by applicable law or agreed to in writing, software 
 * distributed under the License is distributed on an "AS IS" BASIS, 
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
 * See the License for the specific language governing permissions and 
 * limitations under the License. 
 */

package tapioca.util;

import java.util.ConcurrentModificationException;
import java.util.List;

import org.easymock.EasyMock;
import org.easymock.IMocksControl;

import tapioca.testing.AppEngineTestCase;
import tapioca.util.EntityGroupBatchJob;

import com.google.appengine.api.datastore.DatastoreFailureException;
import com.google.appengine.api.datastore.DatastoreService;
import com.google.appengine.api.datastore.DatastoreTimeoutException;
import com.google.appengine.api.datastore.Entity;
import com.google.appengine.api.datastore.EntityNotFoundException;
import com.google.appengine.api.datastore.Key;
import com.google.appengine.api.datastore.Transaction;
import com.google.appengine.repackaged.com.google.common.collect.Lists;

/**
 * Tests for the {@link EntityGroupBatchJob} class.
 *
 * @author armbrust@gmail.com (Erick Armbrust)
 */
public class EntityGroupBatchJobTest extends AppEngineTestCase {

  private final IMocksControl mockControl = EasyMock.createControl();
  private final DatastoreService mockDs = mockControl.createMock(DatastoreService.class);
  private final Transaction mockTxn = mockControl.createMock(Transaction.class);
  
  @Override 
  public void setUp() throws Exception {
    super.setUp();
    mockControl.reset();
  }

  @Override
  public void tearDown() throws Exception {
    mockControl.verify();
    super.tearDown();
  }
    
  private void retryableFailureTest(RuntimeException toThrow) throws Exception {
    Entity f1Eg1 = new Entity("foo", "bar");
    Entity f2Eg1 = new Entity("foo", "bar", f1Eg1.getKey());
    Entity f3Eg2 = new Entity("foo", "baz");
    Entity f4Eg2 = new Entity("foo", "baz", f3Eg2.getKey());
    List<Key> entities = Lists.newArrayList(
        f1Eg1.getKey(), 
        f2Eg1.getKey(), 
        f3Eg2.getKey(), 
        f4Eg2.getKey());
    
    // First entity group will complete successfully.
    EasyMock.expect(mockDs.beginTransaction()).andReturn(mockTxn);
    EasyMock.expect(mockDs.get(mockTxn, f1Eg1.getKey())).andReturn(f1Eg1);
    mockDs.delete(mockTxn, f1Eg1.getKey());
    EasyMock.expect(mockDs.get(mockTxn, f2Eg1.getKey())).andReturn(f2Eg1);
    mockDs.delete(mockTxn, f2Eg1.getKey());
    EasyMock.expect(mockTxn.isActive()).andReturn(true);
    mockTxn.commit();
    
    // First try for this entity group should fail on commit, forcing a retry.
    EasyMock.expect(mockDs.beginTransaction()).andReturn(mockTxn);
    EasyMock.expect(mockDs.get(mockTxn, f3Eg2.getKey())).andReturn(f3Eg2);
    mockDs.delete(mockTxn, f3Eg2.getKey());
    EasyMock.expect(mockDs.get(mockTxn, f4Eg2.getKey())).andReturn(f4Eg2);
    mockDs.delete(mockTxn, f4Eg2.getKey());
    EasyMock.expect(mockTxn.isActive()).andReturn(true);
    mockTxn.commit(); 
    EasyMock.expectLastCall().andThrow(toThrow);
    EasyMock.expect(mockTxn.isActive()).andReturn(false);
    EasyMock.expect(mockDs.beginTransaction()).andReturn(mockTxn);

    // This retry should succeed.
    EasyMock.expect(mockDs.get(mockTxn, f3Eg2.getKey())).andReturn(f3Eg2);
    mockDs.delete(mockTxn, f3Eg2.getKey());
    EasyMock.expect(mockDs.get(mockTxn, f4Eg2.getKey())).andReturn(f4Eg2);
    mockDs.delete(mockTxn, f4Eg2.getKey());
    EasyMock.expect(mockTxn.isActive()).andReturn(true);
    mockTxn.commit(); 
    EasyMock.expectLastCall();
    EasyMock.expect(mockTxn.isActive()).andReturn(false);
    
    mockControl.replay();
    
    EntityGroupBatchJob<Key> batchJob = new EntityGroupBatchJob<Key>(mockDs, entities) {
      @Override
      public Key getKeyFromItem(Key key) {
        return key;
      }

      @Override
      public void processItem(Transaction txn, Key key) {
        try {
          // Check to make sure it is in the datastore.
          mockDs.get(txn, key);
          mockDs.delete(txn, key);
        } catch (EntityNotFoundException e) {
          fail("Unexpected EntityNotFoundException");
        }
      }
    };
    batchJob.execute();        
  }
  
  public void testRunJob() throws Exception {    
    Entity f1Eg1 = new Entity("foo", "bar");
    Entity f2Eg1 = new Entity("foo", "bar", f1Eg1.getKey());
    Entity f3Eg2 = new Entity("foo", "baz");
    Entity f4Eg2 = new Entity("foo", "baz", f3Eg2.getKey());
    List<Key> entities = Lists.newArrayList(
        f1Eg1.getKey(), 
        f2Eg1.getKey(), 
        f3Eg2.getKey(), 
        f4Eg2.getKey());
    
    EasyMock.expect(mockDs.beginTransaction()).andReturn(mockTxn).times(2);
    EasyMock.expect(mockDs.get(mockTxn, f1Eg1.getKey())).andReturn(f1Eg1);
    EasyMock.expect(mockDs.put(mockTxn, f1Eg1)).andReturn(f1Eg1.getKey());
    EasyMock.expect(mockDs.get(mockTxn, f2Eg1.getKey())).andReturn(f2Eg1);
    EasyMock.expect(mockDs.put(mockTxn, f2Eg1)).andReturn(f2Eg1.getKey());
    EasyMock.expect(mockDs.get(mockTxn, f3Eg2.getKey())).andReturn(f3Eg2);
    EasyMock.expect(mockDs.put(mockTxn, f3Eg2)).andReturn(f3Eg2.getKey());
    EasyMock.expect(mockDs.get(mockTxn, f4Eg2.getKey())).andReturn(f4Eg2);
    EasyMock.expect(mockDs.put(mockTxn, f4Eg2)).andReturn(f4Eg2.getKey());
    EasyMock.expect(mockTxn.isActive()).andReturn(true).times(2);
    EasyMock.expect(mockTxn.isActive()).andReturn(false);
    mockTxn.commit();
    EasyMock.expectLastCall().times(2);
    
    mockControl.replay();
    
    EntityGroupBatchJob<Key> batchJob = new EntityGroupBatchJob<Key>(mockDs, entities) {
      @Override
      public Key getKeyFromItem(Key key) {
        return key;
      }

      @Override
      public void processItem(Transaction txn, Key key) {
        try {
          Entity entity = mockDs.get(txn, key);
          mockDs.put(txn, entity);
        } catch (EntityNotFoundException e) {
          fail("Unexpected EntityNotFoundException");
        }
      }
    };
    batchJob.execute();
  }

  public void testRunJob_datastoreConcurrentModificationException() throws Exception {
    retryableFailureTest(new ConcurrentModificationException("fail"));
  }
  
  public void testRunJob_datastoreTimeoutException() throws Exception {
    retryableFailureTest(new DatastoreTimeoutException("fail"));
  }  
  
  public void testRunJob_datastoreFailureException() throws Exception {
    retryableFailureTest(new DatastoreFailureException("fail"));
  }
  
  public void testRunJob_unrecognizedException() throws Exception {
    Entity f1Eg1 = new Entity("foo", "bar");
    Entity f2Eg1 = new Entity("foo", "bar", f1Eg1.getKey());
    Entity f3Eg2 = new Entity("foo", "baz");
    Entity f4Eg2 = new Entity("foo", "baz", f3Eg2.getKey());
    List<Key> entities = Lists.newArrayList(
        f1Eg1.getKey(), 
        f2Eg1.getKey(), 
        f3Eg2.getKey(), 
        f4Eg2.getKey());
    RuntimeException fake = new RuntimeException("fake");
    
    // First entity group will complete successfully.
    EasyMock.expect(mockDs.beginTransaction()).andReturn(mockTxn);
    EasyMock.expect(mockDs.get(mockTxn, f1Eg1.getKey())).andReturn(f1Eg1);
    mockDs.delete(mockTxn, f1Eg1.getKey());
    EasyMock.expect(mockDs.get(mockTxn, f2Eg1.getKey())).andReturn(f2Eg1);
    mockDs.delete(mockTxn, f2Eg1.getKey());
    EasyMock.expect(mockTxn.isActive()).andReturn(true);
    mockTxn.commit();
    
    // First try for this entity group should fail on commit, forcing a retry.
    EasyMock.expect(mockDs.beginTransaction()).andReturn(mockTxn);
    EasyMock.expect(mockDs.get(mockTxn, f3Eg2.getKey())).andReturn(f3Eg2);
    mockDs.delete(mockTxn, f3Eg2.getKey());
    EasyMock.expect(mockDs.get(mockTxn, f4Eg2.getKey())).andReturn(f4Eg2);
    mockDs.delete(mockTxn, f4Eg2.getKey());
    EasyMock.expect(mockTxn.isActive()).andReturn(true);
    mockTxn.commit(); 
    EasyMock.expectLastCall().andThrow(fake);
    EasyMock.expect(mockTxn.isActive()).andReturn(false);
    
    mockControl.replay();
    
    EntityGroupBatchJob<Key> batchJob = new EntityGroupBatchJob<Key>(mockDs, entities) {
      @Override
      public Key getKeyFromItem(Key key) {
        return key;
      }

      @Override
      public void processItem(Transaction txn, Key key) {
        try {
          // Check to make sure it is in the datastore.
          mockDs.get(txn, key);
          mockDs.delete(txn, key);
        } catch (EntityNotFoundException e) {
          fail("Unexpected EntityNotFoundException");
        }
      }
    };
    
    try {
      batchJob.execute();
      fail("Unknown exception should have caused this DT to exit.");
    } catch (RuntimeException e) {
      assertEquals(fake, e);
    }
  }
  
}
